// Copyright (c) 2022-present, Shanghai Yunxi Technology Co, Ltd.
//
// This software (KWDB) is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//          http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
// EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

#include "ee_sort_op.h"

#include <queue>
#include <utility>
#include <memory>

#include "cm_func.h"
#include "ee_kwthd_context.h"
#include "ee_pb_plan.pb.h"
#include "lg_api.h"

namespace kwdbts {

SortOperator::SortOperator(TsFetcherCollection* collection, TSSorterSpec* spec,
                           PostProcessSpec* post, TABLE* table, int32_t processor_id)
    : BaseOperator(collection, table, post, processor_id),
      spec_{spec},
      param_(spec, post, table),
      limit_(post->limit()),
      offset_(post->offset()) {}

SortOperator::SortOperator(const SortOperator& other, int32_t processor_id)
    : BaseOperator(other),
      spec_(other.spec_),
      param_(other.spec_, other.post_, other.table_),
      limit_(other.post_->limit()),
      offset_(other.post_->offset()) {
  is_clone_ = true;
}

SortOperator::~SortOperator() {
  //  delete input_
  if (is_clone_) {
    delete childrens_[0];
  }
  SafeDeleteArray(input_col_info_);
}

KStatus SortOperator::ResolveSortCols(kwdbContext_p ctx) {
  EnterFunc();
  if (!spec_->has_output_ordering()) {
    LOG_ERROR("order by clause must has a order field");
    EEPgErrorInfo::SetPgErrorInfo(ERRCODE_INVALID_PARAMETER_VALUE,
                                  "order by clause must has a order field");
    Return(KStatus::FAIL);
  }

  const TSOrdering& ordering = spec_->output_ordering();
  int order_size_ = ordering.columns_size();

  for (k_int32 i = 0; i < order_size_; i++) {
    k_uint32 idx = ordering.columns(i).col_idx();
    TSOrdering_Column_Direction direction = ordering.columns(i).direction();
    order_info_.push_back(ColumnOrderInfo{idx, direction});
  }

  Return(KStatus::SUCCESS);
}

EEIteratorErrCode SortOperator::Init(kwdbContext_p ctx) {
  EnterFunc();
  EEIteratorErrCode code = EEIteratorErrCode::EE_ERROR;

  do {
    param_.AddInput(childrens_[0]);
    std::vector<Field *> &input_fields = childrens_[0]->OutputFields();
    // input preinit
    code = childrens_[0]->Init(ctx);
    if (EEIteratorErrCode::EE_OK != code) {
      break;
    }
    // resolve renders num
    param_.RenderSize(ctx, &num_);
    if (0 != num_) {
      // resolve render
      code = param_.ParserRender(ctx, &renders_, num_);
      if (EEIteratorErrCode::EE_OK != code) {
        LOG_ERROR("ResolveRender() error\n");
        break;
      }
    } else {
      k_uint32 num = input_fields.size();
      renders_ = static_cast<Field **>(malloc(num * sizeof(Field *)));
      if (!renders_) {
        EEPgErrorInfo::SetPgErrorInfo(ERRCODE_OUT_OF_MEMORY,
                                      "Insufficient memory");
        LOG_ERROR("Malloc failed, size : %lu", num * sizeof(Field*));
        break;
      }
      num_ = num;
      for (k_uint32 i = 0; i < num_; i++) {
        renders_[i] = input_fields[i];
      }
    }
    // dispose sort col
    KStatus ret = ResolveSortCols(ctx);
    if (ret != KStatus::SUCCESS) {
      code = EEIteratorErrCode::EE_ERROR;
      break;
    }

    // output Field
    code = param_.ParserOutputFields(ctx, renders_, num_, output_fields_, false);
    if (code != EEIteratorErrCode::EE_OK) {
      break;
    }
    code = InitOutputColInfo(output_fields_);
    if (code != EEIteratorErrCode::EE_OK) {
      break;
    }
    input_col_num_ = input_fields.size();
    input_col_info_ = KNEW ColumnInfo[input_col_num_];
    if (input_col_info_ == nullptr) {
      EEPgErrorInfo::SetPgErrorInfo(ERRCODE_OUT_OF_MEMORY,
                                    "Insufficient memory");
      Return(EEIteratorErrCode::EE_ERROR);
    }
    for (k_int32 i = 0; i < input_fields.size(); i++) {
      input_col_info_[i] = ColumnInfo(input_fields[i]->get_storage_length(),
                                       input_fields[i]->get_storage_type(),
                                       input_fields[i]->get_return_type());
    }
  } while (0);

  Return(code);
}

EEIteratorErrCode SortOperator::Start(kwdbContext_p ctx) {
  EnterFunc();
  EEIteratorErrCode code = EEIteratorErrCode::EE_ERROR;
  cur_offset_ = offset_;

  code = childrens_[0]->Start(ctx);
  if (EEIteratorErrCode::EE_OK != code) {
    Return(code);
  }

  std::queue<DataChunkPtr> buffer;
  size_t total_count = 0;
  size_t buffer_size = 0;
  int64_t duration = 0;
  // init container
  container_ =
      std::make_unique<MemRowContainer>(order_info_, input_col_info_,
                                        input_col_num_);
  KStatus ret = container_->Init();
  if (ret != SUCCESS) {
    container_ = nullptr;
    EEPgErrorInfo::SetPgErrorInfo(ERRCODE_INTERNAL_ERROR,
                                  "Init sort container.");
    Return(EEIteratorErrCode::EE_ERROR);
  }
  if (limit_ > 0) {
    container_->SetLimitOffset(limit_, offset_);
  }
  // sort all data
  for (;;) {
    DataChunkPtr chunk = nullptr;

    // read a batch of data
    code = childrens_[0]->Next(ctx, chunk);
    auto start = std::chrono::high_resolution_clock::now();
    if (code != EEIteratorErrCode::EE_OK) {
      if (code == EEIteratorErrCode::EE_END_OF_RECORD ||
          code == EEIteratorErrCode::EE_TIMESLICE_OUT) {
        code = EEIteratorErrCode::EE_OK;
        break;
      }
      // LOG_ERROR("Failed to fetch data from child operator, return code = %d.\n",
      //           code);
      Return(code);
    }
    // no data, continue
    // if (chunk == nullptr || chunk->Count() == 0) {
    //   continue;
    // }

    total_count += chunk->Count();
    buffer_size += chunk->RowSize() * chunk->Count();
    KStatus ret = SUCCESS;
    if (sort_type_ == EESortType::EE_SORT_MEMORY) {
      if (buffer_size > SORT_MAX_MEM_BUFFER_SIZE) {
        sort_type_ = EESortType::EE_SORT_DISK;
        auto disk_container = std::make_unique<DiskDataContainer>(
            order_info_, input_col_info_, input_col_num_);
        ret = disk_container->Init();
        if (ret != SUCCESS) {
          Return(EEIteratorErrCode::EE_ERROR);
        }
        if (limit_ > 0) {
          disk_container->SetLimitOffset(limit_, offset_);
        }
        while (true) {
          DataChunkPtr mem_chunk;
          auto code = container_->NextChunk(mem_chunk);
          if (code != EEIteratorErrCode::EE_OK) {
            break;
          }
          ret = disk_container->Append(mem_chunk);
          if (ret != SUCCESS) {
            Return(EEIteratorErrCode::EE_ERROR);
          }
        }
        container_ = std::move(disk_container);
      } else if (limit_ > 0 && (limit_ + offset_) <
                                   (SORT_MAX_MEM_BUFFER_SIZE / chunk->RowSize())) {
        sort_type_ = EESortType::EE_SORT_HEAP;
        auto heap_container = std::make_unique<HeapSortContainer>(
            order_info_, input_col_info_, input_col_num_, limit_ + offset_);
        ret = heap_container->Init();
        if (ret != SUCCESS) {
          Return(EEIteratorErrCode::EE_ERROR);
        }
        heap_container->SetLimitOffset(limit_, offset_);
        while (true) {
          DataChunkPtr mem_chunk;
          auto code = container_->NextChunk(mem_chunk);
          if (code != EEIteratorErrCode::EE_OK) {
            break;
          }
          ret = heap_container->Append(mem_chunk);
          if (ret != SUCCESS) {
            Return(EEIteratorErrCode::EE_ERROR);
          }
        }
        container_ = std::move(heap_container);
      }
    }
    ret = container_->Append(chunk);
    // }
    if (ret != SUCCESS) {
      EEPgErrorInfo::SetPgErrorInfo(ERRCODE_INTERNAL_ERROR,
                                    "Append data failed.");
      Return(EEIteratorErrCode::EE_ERROR);
    }
    auto end = std::chrono::high_resolution_clock::now();
    fetcher_.Update(0, (end - start).count(), 0, 0, 0, 0);
  }

  auto start = std::chrono::high_resolution_clock::now();
  // Sort
  ret = container_->Sort();
  if (ret != SUCCESS) {
    EEPgErrorInfo::SetPgErrorInfo(ERRCODE_INTERNAL_ERROR, "Sort data failed.");
    Return(EEIteratorErrCode::EE_ERROR);
  }
  auto end = std::chrono::high_resolution_clock::now();
  fetcher_.Update(total_count, (end - start).count(), 0, buffer_size, 0, 0);
  Return(code);
}

EEIteratorErrCode SortOperator::Next(kwdbContext_p ctx,
                                         DataChunkPtr& chunk) {
  EnterFunc();
  EEIteratorErrCode code = EEIteratorErrCode::EE_ERROR;
  KWThdContext* thd = current_thd;
  if (is_done_) {
    Return(EEIteratorErrCode::EE_END_OF_RECORD);
  }
  auto start = std::chrono::high_resolution_clock::now();
  if (nullptr == chunk) {
    chunk = std::make_unique<DataChunk>(output_col_info_, output_col_num_);
    if (chunk->Initialize() != true) {
      EEPgErrorInfo::SetPgErrorInfo(ERRCODE_OUT_OF_MEMORY,
                                    "Insufficient memory");
      chunk = nullptr;
      Return(EEIteratorErrCode::EE_ERROR);
    }
  }
  k_uint32 BATCH_SIZE = chunk->Capacity();
    thd->SetDataChunk(container_.get());
    // record location in current result batch.
    k_uint32 location = 0;
    while (scanned_rows_ < container_->Count()) {
      // limit
      if (limit_ && examined_rows_ >= limit_) {
        is_done_ = true;
        break;
      }
      k_int32 row = container_->NextLine();
      if (row < 0) {
        is_done_ = true;
        break;
      }
      ++scanned_rows_;

      // offset
      if (cur_offset_ > 0) {
        --cur_offset_;
        continue;
      }
      chunk->InsertData(ctx, container_.get(), num_ != 0 ? renders_ : nullptr);

      // rowcount ++
      ++examined_rows_;
      ++location;

      if (examined_rows_ % BATCH_SIZE == 0) {
        break;
      }
    }

    if (scanned_rows_ == container_->Count()) {
      is_done_ = true;
    }

  if (0 == chunk->Count()) {
    chunk = nullptr;
    Return(EEIteratorErrCode::EE_END_OF_RECORD);
  }

  OPERATOR_DIRECT_ENCODING(ctx, output_encoding_, use_query_short_circuit_, thd, chunk);
  auto end = std::chrono::high_resolution_clock::now();
  fetcher_.Update(0, (end - start).count(), chunk->Count() * chunk->RowSize(),
                  0, 0, 0);
  Return(EEIteratorErrCode::EE_OK);
}

EEIteratorErrCode SortOperator::Close(kwdbContext_p ctx) {
  EnterFunc();
  EEIteratorErrCode code = childrens_[0]->Close(ctx);
  Reset(ctx);

  Return(code);
}

EEIteratorErrCode SortOperator::Reset(kwdbContext_p ctx) {
  EnterFunc();
  childrens_[0]->Reset(ctx);

  Return(EEIteratorErrCode::EE_OK);
}

BaseOperator* SortOperator::Clone() {
  BaseOperator* input = childrens_[0]->Clone();
  if (input == nullptr) {
    return nullptr;
  }
  BaseOperator* iter = NewIterator<SortOperator>(*this, this->processor_id_);
  if (nullptr != iter) {
    iter->AddDependency(input);
  } else {
    delete input;
  }

  return iter;
}

}  // namespace kwdbts
